from llvmlite import binding as llvm
from ctypes import Structure, c_int32, c_double, POINTER, CFUNCTYPE, byref

# 初始化 LLVM
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()

# 定义 LLVM 汇编代码（使用指针传递返回值）
ir_text = """
%MyStruct = type { i32, double }

; 通过指针参数返回结构体的函数
define void @create_struct_ptr(%MyStruct* noalias sret(%MyStruct) %result) {
  ; 设置结构体字段
  %a_ptr = getelementptr inbounds %MyStruct, %MyStruct* %result, i32 0, i32 0
  store i32 42, i32* %a_ptr
  
  %b_ptr = getelementptr inbounds %MyStruct, %MyStruct* %result, i32 0, i32 1
  store double 3.140000e+00, double* %b_ptr
  
  ret void
}

; 通过指针参数返回多个值的函数
define void @calculate_values(i32 %x, double %y, %MyStruct* %result) {
  ; 计算新值
  %sum = add i32 %x, 100
  %product = fmul double %y, 2.0
  
  ; 存储到结果结构体
  %a_ptr = getelementptr inbounds %MyStruct, %MyStruct* %result, i32 0, i32 0
  store i32 %sum, i32* %a_ptr
  
  %b_ptr = getelementptr inbounds %MyStruct, %MyStruct* %result, i32 0, i32 1
  store double %product, double* %b_ptr
  
  ret void
}

; 创建并返回结构体指针的函数
define %MyStruct* @create_and_return_ptr() {
  ; 分配内存
  %ptr = alloca %MyStruct
  
  ; 初始化结构体
  %a_ptr = getelementptr inbounds %MyStruct, %MyStruct* %ptr, i32 0, i32 0
  store i32 200, i32* %a_ptr
  
  %b_ptr = getelementptr inbounds %MyStruct, %MyStruct* %ptr, i32 0, i32 1
  store double 6.280000e+00, double* %b_ptr
  
  ret %MyStruct* %ptr
}
"""

# 解析汇编代码
module = llvm.parse_assembly(ir_text)
module.verify()

# 创建执行引擎
target = llvm.Target.from_default_triple()
target_machine = target.create_target_machine()
engine = llvm.create_mcjit_compiler(module, target_machine)

# 定义匹配的 Python 结构体
class MyStruct(Structure):
    _fields_ = [("a", c_int32), ("b", c_double)]
    
    def __str__(self):
        return f"MyStruct(a={self.a}, b={self.b})"

# 获取函数地址并设置调用签名
# 1. 通过指针参数返回结构体
create_struct_ptr_addr = engine.get_function_address("create_struct_ptr")
create_struct_ptr_func = CFUNCTYPE(None, POINTER(MyStruct))(create_struct_ptr_addr)

# 2. 通过指针参数返回计算结果
calculate_values_addr = engine.get_function_address("calculate_values")
calculate_values_func = CFUNCTYPE(None, c_int32, c_double, POINTER(MyStruct))(calculate_values_addr)

# 3. 返回结构体指针的函数
create_and_return_ptr_addr = engine.get_function_address("create_and_return_ptr")
create_and_return_ptr_func = CFUNCTYPE(POINTER(MyStruct))(create_and_return_ptr_addr)

# 测试函数调用
if __name__ == "__main__":
    print("=== 通过指针参数返回结构体 ===")
    result1 = MyStruct()
    create_struct_ptr_func(byref(result1))
    print(f"Result: {result1}")  # 输出: MyStruct(a=42, b=3.14)
    
    print("\n=== 通过指针参数返回计算结果 ===")
    result2 = MyStruct()
    calculate_values_func(50, 1.5, byref(result2))
    print(f"Result: {result2}")  # 输出: MyStruct(a=150, b=3.0)
    
    print("\n=== 返回结构体指针 ===")
    ptr = create_and_return_ptr_func()
    print(f"Pointer: {ptr}")
    print(f"Content: {ptr.contents}")  # 输出: MyStruct(a=200, b=6.28)
    
    # 注意：返回的指针指向 LLVM 分配的内存
    # 在实际应用中需要考虑内存管理
